import math
from typing import Optional

import torch
import torch.nn.functional as F
from ..utils import pytorch_toolbelt_deprecated

__all__ = [
    "focal_loss_with_logits",
    "softmax_focal_loss_with_logits",
    "sigmoid_focal_loss",
    "soft_jaccard_score",
    "soft_dice_score",
    "wing_loss",
    "log_cosh_loss",
]


def focal_loss_with_logits(
    output: torch.Tensor,
    target: torch.Tensor,
    gamma: float = 2.0,
    alpha: Optional[float] = 0.25,
    reduction: str = "mean",
    normalized: bool = False,
    reduced_threshold: Optional[float] = None,
    eps: float = 1e-6,
    ignore_index=None,
    activation: str = "sigmoid",
    softmax_dim: Optional[int] = None,
    class_weights: Optional[torch.Tensor] = None,
) -> torch.Tensor:
    """Compute binary focal loss between target and output logits.

    See :class:`~pytorch_toolbelt.losses.FocalLoss` for details.

    Args:
        output: Tensor of arbitrary shape (predictions of the model)
        target: Tensor of the same shape as input
        gamma: Focal loss power factor
        alpha: Weight factor to balance positive and negative samples. Alpha must be in [0...1] range,
            high values will give more weight to positive class.
        reduction (string, optional): Specifies the reduction to apply to the output:
            'none' | 'mean' | 'sum' | 'batchwise_mean'. 'none': no reduction will be applied,
            'mean': the sum of the output will be divided by the number of
            elements in the output, 'sum': the output will be summed. Note: :attr:`size_average`
            and :attr:`reduce` are in the process of being deprecated, and in the meantime,
            specifying either of those two args will override :attr:`reduction`.
            'batchwise_mean' computes mean loss per sample in batch. Default: 'mean'
        normalized (bool): Compute normalized focal loss (https://arxiv.org/pdf/1909.07829.pdf).
        reduced_threshold (float, optional): Compute reduced focal loss (https://arxiv.org/abs/1903.01347).
        activation: Either sigmoid or softmax. If `softmax` is used, `softmax_dim` must be also specified.

    References:
        https://github.com/open-mmlab/mmdetection/blob/master/mmdet/core/loss/losses.py
    """
    output = output.float()
    target = target.float()

    with torch.amp.autocast(device_type=output.device.type, enabled=False):
        if activation == "sigmoid":
            p = torch.sigmoid(output)
        else:
            p = torch.softmax(output, dim=softmax_dim)

        ce_loss = F.binary_cross_entropy_with_logits(output, target, reduction="none")
        pt = p * target + (1 - p) * (1 - target)

        # compute the loss
        if reduced_threshold is None:
            focal_term = (1.0 - pt).pow(gamma)
        else:
            focal_term = ((1.0 - pt) / (1 - reduced_threshold)).pow(
                gamma
            )  # the focal term continuity breaks when reduced_threshold not equal to 0.5. At pt equal to reduced_threshold, the value of piecewise function of focal term should be 1 from both sides .
            focal_term = torch.masked_fill(focal_term, pt < reduced_threshold, 1)

        loss = focal_term * ce_loss

        if alpha is not None:
            loss *= alpha * target + (1 - alpha) * (1 - target)

        if class_weights is not None:
            # class_weights is of shape [C]
            # Loss is of shape [B,C ...]
            # Reshape class_weights to [1, C, ...]
            class_weights = class_weights.view(1, -1, *(1 for _ in range(loss.dim() - 2)))
            loss *= class_weights

        if ignore_index is not None:
            ignore_mask = target.eq(ignore_index)
            loss = torch.masked_fill(loss, ignore_mask, 0)
            if normalized:
                focal_term = torch.masked_fill(focal_term, ignore_mask, 0)

        if normalized:
            norm_factor = focal_term.sum(dtype=torch.float32).clamp_min(eps)
            loss /= norm_factor

        if reduction == "mean":
            loss = loss.mean()
        if reduction == "sum":
            loss = loss.sum()
        if reduction == "batchwise_mean":
            loss = loss.sum(dim=0)

    return loss


def softmax_focal_loss_with_logits(
    output: torch.Tensor,
    target: torch.Tensor,
    class_weights: Optional[torch.Tensor] = None,
    gamma: float = 2.0,
    reduction: str = "mean",
    normalized: bool = False,
    reduced_threshold: Optional[float] = None,
    eps: float = 1e-6,
    ignore_index: int = -100,
) -> torch.Tensor:
    """
    Softmax version of focal loss between target and output logits.
    See :class:`~pytorch_toolbelt.losses.CrossEntropyFocalLoss` for details.

    Args:
        output: Tensor of shape [B, C, *] (Similar to nn.CrossEntropyLoss)
        target: Tensor of shape [B, *] (Similar to nn.CrossEntropyLoss)
        gamma: Focal loss power factor
        reduction (string, optional): Specifies the reduction to apply to the output:
            'none' | 'mean' | 'sum' | 'batchwise_mean'. 'none': no reduction will be applied,
            'mean': the sum of the output will be divided by the number of
            elements in the output, 'sum': the output will be summed. Note: :attr:`size_average`
            and :attr:`reduce` are in the process of being deprecated, and in the meantime,
            specifying either of those two args will override :attr:`reduction`.
            'batchwise_mean' computes mean loss per sample in batch. Default: 'mean'
        normalized (bool): Compute normalized focal loss (https://arxiv.org/pdf/1909.07829.pdf).
        reduced_threshold (float, optional): Compute reduced focal loss (https://arxiv.org/abs/1903.01347).
    """
    ignore_mask = target.eq(ignore_index)
    pos_mask = ~ignore_mask
    targets_masked = torch.masked_fill(target, ignore_mask, 0)
    targets_oh = F.one_hot(targets_masked, num_classes=output.size(1)).moveaxis(-1, 1).float()
    probs = F.softmax(output, dim=1)
    pt = (1 - targets_oh) * probs + targets_oh * (1 - probs)

    loss = F.binary_cross_entropy_with_logits(output, targets_oh, reduction="none")

    # compute the loss
    if reduced_threshold is None:
        focal_term = pt.pow(gamma)
    else:
        focal_term = (pt / reduced_threshold).pow(gamma)
        focal_term[pt < reduced_threshold] = 1

    loss = focal_term * loss
    if class_weights is not None:
        class_weights = class_weights.reshape(1, -1, *[1 for _ in range(loss.dim() - 2)])
        loss = loss * class_weights
    loss = loss.sum(dim=1) * pos_mask

    if normalized:
        # Does not really work, will be removed in future
        norm_factor = focal_term.sum().clamp_min(eps)
        loss = loss / norm_factor

    if reduction == "mean":
        loss = loss.mean()
    if reduction == "sum":
        loss = loss.sum()
    if reduction == "batchwise_mean":
        loss = loss.sum(0)

    return loss


@pytorch_toolbelt_deprecated("Function sigmoid_focal_loss is deprecated. Please use focal_loss_with_logits instead.")
def sigmoid_focal_loss(*input, **kwargs):
    return focal_loss_with_logits(*input, **kwargs)


@pytorch_toolbelt_deprecated("Function reduced_focal_loss is deprecated. Please use focal_loss_with_logits instead.")
def reduced_focal_loss(output: torch.Tensor, target: torch.Tensor, threshold=0.5, gamma=2.0, reduction="mean"):
    return focal_loss_with_logits(
        output, target, alpha=None, gamma=gamma, reduction=reduction, reduced_threshold=threshold
    )


def soft_jaccard_score(
    output: torch.Tensor, target: torch.Tensor, smooth: float = 0.0, eps: float = 1e-7, dims=None
) -> torch.Tensor:
    """

    :param output:
    :param target:
    :param smooth:
    :param eps:
    :param dims:
    :return:

    Shape:
        - Input: :math:`(N, NC, *)` where :math:`*` means
            any number of additional dimensions
        - Target: :math:`(N, NC, *)`, same shape as the input
        - Output: scalar.

    """
    assert output.size() == target.size()

    if dims is not None:
        intersection = torch.sum(output * target, dim=dims)
        cardinality = torch.sum(output + target, dim=dims)
    else:
        intersection = torch.sum(output * target)
        cardinality = torch.sum(output + target)

    union = cardinality - intersection
    jaccard_score = (intersection + smooth) / (union + smooth).clamp_min(eps)
    return jaccard_score


def soft_dice_score(
    output: torch.Tensor, target: torch.Tensor, smooth: float = 0.0, eps: float = 1e-7, dims=None
) -> torch.Tensor:
    """

    :param output:
    :param target:
    :param smooth:
    :param eps:
    :return:

    Shape:
        - Input: :math:`(N, NC, *)` where :math:`*` means any number
            of additional dimensions
        - Target: :math:`(N, NC, *)`, same shape as the input
        - Output: scalar.

    """
    assert output.size() == target.size()
    if dims is not None:
        intersection = torch.sum(output * target, dim=dims)
        cardinality = torch.sum(output + target, dim=dims)
    else:
        intersection = torch.sum(output * target)
        cardinality = torch.sum(output + target)
    dice_score = (2.0 * intersection + smooth) / (cardinality + smooth).clamp_min(eps)
    return dice_score


def wing_loss(output: torch.Tensor, target: torch.Tensor, width=5, curvature=0.5, reduction="mean"):
    """
    https://arxiv.org/pdf/1711.06753.pdf
    :param output:
    :param target:
    :param width:
    :param curvature:
    :param reduction:
    :return:
    """
    diff_abs = (target - output).abs()
    loss = diff_abs.clone()

    idx_smaller = diff_abs < width
    idx_bigger = diff_abs >= width

    loss[idx_smaller] = width * torch.log(1 + diff_abs[idx_smaller] / curvature)

    C = width - width * math.log(1 + width / curvature)
    loss[idx_bigger] = loss[idx_bigger] - C

    if reduction == "sum":
        loss = loss.sum()

    if reduction == "mean":
        loss = loss.mean()

    return loss


def label_smoothed_nll_loss(
    lprobs: torch.Tensor, target: torch.Tensor, epsilon: float, ignore_index=None, reduction="mean", dim=-1
) -> torch.Tensor:
    """

    Source: https://github.com/pytorch/fairseq/blob/master/fairseq/criterions/label_smoothed_cross_entropy.py

    :param lprobs: Log-probabilities of predictions (e.g after log_softmax)
    :param target:
    :param epsilon:
    :param ignore_index:
    :param reduction:
    :return:
    """
    if target.dim() == lprobs.dim() - 1:
        target = target.unsqueeze(dim)

    if ignore_index is not None:
        pad_mask = target.eq(ignore_index)
        target = target.masked_fill(pad_mask, 0)
        nll_loss = -lprobs.gather(dim=dim, index=target)
        smooth_loss = -lprobs.sum(dim=dim, keepdim=True)

        # nll_loss.masked_fill_(pad_mask, 0.0)
        # smooth_loss.masked_fill_(pad_mask, 0.0)
        nll_loss = nll_loss.masked_fill(pad_mask, 0.0)
        smooth_loss = smooth_loss.masked_fill(pad_mask, 0.0)
    else:
        nll_loss = -lprobs.gather(dim=dim, index=target)
        smooth_loss = -lprobs.sum(dim=dim, keepdim=True)

        nll_loss = nll_loss.squeeze(dim)
        smooth_loss = smooth_loss.squeeze(dim)

    if reduction == "sum":
        nll_loss = nll_loss.sum()
        smooth_loss = smooth_loss.sum()
    if reduction == "mean":
        nll_loss = nll_loss.mean()
        smooth_loss = smooth_loss.mean()

    eps_i = epsilon / lprobs.size(dim)
    loss = (1.0 - epsilon) * nll_loss + eps_i * smooth_loss
    return loss


def log_cosh_loss(y_pred: torch.Tensor, y_true: torch.Tensor) -> torch.Tensor:
    """
    Numerically stable log-cosh implementation.
    Reference: https://datascience.stackexchange.com/questions/96271/logcoshloss-on-pytorch

    Args:
        y_pred:
        y_true:

    Returns:

    """

    def _log_cosh(x: torch.Tensor) -> torch.Tensor:
        return x + torch.nn.functional.softplus(-2.0 * x) - math.log(2.0)

    return torch.mean(_log_cosh(y_pred - y_true))
